/*
 * This is written by Lulu Han.
 * E-mail: locomotive_crypto@163.com
 */

#include <iostream>
#include <vector>
#include <exception>
#include "paillier.h"
#include "timer.h"
#include "net_client.h"
#include "net_server.h"

/*Testing for Paillier encryption*/

void test_for_correctness();

void test_paillier_for_performance(long bits, long cnt);

void test_paillier_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr);

void test_paillier_by_network_SB(size_t port);

int main(int argc, char* argv[]){
	std::cout << "----------------------------------------------------------" << std::endl;
	std::cout << "-----------------------<Usage>----------------------------" << std::endl;
	std::cout << "-----1) Run for correctness: ./program--------------------" << std::endl;
	std::cout << "-----2) Run for performance: ./program cnt----------------" << std::endl;
	std::cout << "----------------------------------------------------------" << std::endl << std::endl;
	
	if(argc == 1){
		test_for_correctness();
	}
	else if(argc == 2){
		for(size_t i = 0; i < 10; ++i){
			test_paillier_for_performance(512, std::atoi(argv[1]));
		
			test_paillier_for_performance(1024, std::atoi(argv[1]));
		
			test_paillier_for_performance(1536, std::atoi(argv[1]));
		}
	}
	else if(argc == 3){
		if(std::string(argv[1]) == std::string("sb") || std::string(argv[1]) == std::string("SB")){
			test_paillier_by_network_SB(std::atoi(argv[2]));
		}
			
	}
	else if(argc == 5){
		if(std::string(argv[1]) == std::string("sa") || std::string(argv[1]) == std::string("SA")){
			//test_paillier_by_network_SA(512, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			test_paillier_by_network_SA(1024, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			//test_paillier_by_network_SA(1536, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
		}
			
	}
	return 0;
}

void test_for_correctness(){
	Paillier::PrivateKey prikey;
	Paillier::PublicKey pubkey;
	
	// Generate random key pair
	Paillier::GenKeyPair(prikey, pubkey, 1024);
	
	// Define Encryptor and Decryptor objects
	Paillier::Encryptor encryptor(pubkey);
	Paillier::Decryptor decryptor(prikey);
	
	// Testing for encryption and decryption
	NTL::ZZ m1(1234567); // plaintext
	NTL::ZZ c1; // ciphertext
	
	encryptor.Encrypt(m1, c1);
	std::cout << "m1 = " << m1 << std::endl;
	
	NTL::ZZ recover;
	decryptor.Decrypt(c1, recover);
	std::cout << "recover(m1) = " << recover << std::endl;
	
	// Testing for homomorphism properties
	NTL::ZZ m2(1234567);
	std::cout << "m2 = " << m2 << std::endl;
	NTL::ZZ c2, c3;
	encryptor.Encrypt(m2, c2); 
	encryptor.HE_Add_Cipher(c1, c2, c3); 
	decryptor.Decrypt(c3, recover);
	std::cout << "recover(m1 + m2) = " << recover << std::endl;
	if(recover != (m1 + m2)){
		throw std::runtime_error("Testing for additive homomorphism failed...");
	}
	
	NTL::ZZ n1(123);
	std::cout << "n1 = " << n1 << std::endl;
	NTL::ZZ c4;
	encryptor.HE_Add_Const(c1, n1, c4); 
	decryptor.Decrypt(c4, recover); 
	std::cout << "recover(m1 + n1) = " << recover << std::endl;
	if(recover != (m1 + n1)){
		throw std::runtime_error("Testing for adding a constant failed...");
	}
	
	NTL::ZZ n2(567);
	std::cout << "n2 = " << n2 << std::endl;
	NTL::ZZ c5;
	encryptor.HE_Mul_Const(c1, n2, c5);
	decryptor.Decrypt(c5, recover); 
	std::cout << "recover(m1 * n2) = " << recover << std::endl;
	if(recover != (m1 * n2)){
		throw std::runtime_error("Testing for scalar multiplication homomorphism failed...");
	}
}

void test_paillier_for_performance(long bits, long cnt){
	
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	 
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<Paillier::Decryptor> vdecryptor;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
	
		// Define Encryptor and Decryptor objects
		Paillier::Encryptor encryptor(pubkey);
		Paillier::Decryptor decryptor(prikey);
		
		vencryptor.push_back(encryptor);
		vdecryptor.push_back(decryptor);
	}
	
	std::vector<NTL::ZZ> vmsg;
	NTL::ZZ r;
	for(long i = 0; i < cnt; ++i){
		r = NTL::RandomBnd(vencryptor[i].GetN());
		vmsg.push_back(r);
	}
	
	time_point<high_resolution_clock> start;	
	start = now_time();
	
	std::vector<NTL::ZZ> vc;
	NTL::ZZ c;
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].Encrypt(vmsg[i], c);
		vc.push_back(c);
	}
	double tm = time_from(start);
	std::cout << "The time for "<< cnt << " encryptions is : " << tm << std::endl;
	
	start = now_time();
	NTL::ZZ m;
	for(long i = 0; i < cnt; ++i){
		vdecryptor[i].Decrypt(vc[i], m);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " decryptions is : " << tm << std::endl;
	
	std::vector<NTL::ZZ> vn;
	for(long i = 0; i < cnt; ++i){
		r = NTL::RandomBnd(vencryptor[i].GetN());
		vn.push_back(r);
	}
	
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].HE_Add_Const(vc[i], vn[i], c);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " HE_Add_Const() is : " << tm << std::endl;
	
	std::vector<NTL::ZZ> vc2;
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].Encrypt(vn[i], c);
		vc2.push_back(c);
	}
	
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].HE_Add_Cipher(vc[i], vc2[i], c);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " HE_Add_Cipher() is : " << tm << std::endl;
	
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].HE_Mul_Const(vc[i], vn[i], c);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " HE_Mul_Const() is : " << tm << std::endl << std::endl;
}

void test_paillier_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	NetClient client(ipaddr, port);
	client.connect();
	
	client.send(cnt);
	client.send(bits);
	
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<Paillier::Decryptor> vdecryptor;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
	
		// Define Encryptor and Decryptor objects
		Paillier::Encryptor encryptor(pubkey);
		Paillier::Decryptor decryptor(prikey);
		
		vencryptor.push_back(encryptor);
		vdecryptor.push_back(decryptor);
	}
	
	std::cout << "The Communication cost for Z_N..." << std::endl;
	getchar();
	
	std::vector<NTL::ZZ> vmsg;
	NTL::ZZ r;
	for(long i = 0; i < cnt; ++i){
		r = NTL::RandomBnd(vencryptor[i].GetN());
		vmsg.push_back(r);
	}
	client.send(vmsg);
	
	std::cout << "The Communication cost for Z_N^2..." << std::endl;
	getchar();
	
	std::vector<NTL::ZZ> vc;
	NTL::ZZ c;
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].Encrypt(vmsg[i], c);
		vc.push_back(c);
	}
	client.send(vc);
	
	client.close();
}

void test_paillier_by_network_SB(size_t port){
	NetServer server(port); 
	server.accept();
	
	long cnt;
	server.receive(cnt);
	
	long bits;
	server.receive(bits);
	
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	std::cout << "The Communication cost for Z_N..." << std::endl;
	getchar();
	std::vector<NTL::ZZ> vmsg;
	server.receive(vmsg);
	
	std::cout << "The Communication cost for Z_N^2..." << std::endl;
	getchar();
	std::vector<NTL::ZZ> vc;
	server.receive(vc);
	
	server.close();
}

